from typing import Optional
import torch
import torch.nn as nn

from torchvision import ops

import numpy as np

from einops import rearrange, reduce, repeat

from mmcv.cnn import build_activation_layer, build_norm_layer

import math

class FourierEmbed(nn.Module):
    """Fourier Embedding Layer from https://arxiv.org/pdf/2106.02795.pdf"""
    def __init__(self, d_embed, n_dim, temperature=16):
        super().__init__()
        assert d_embed % 2 == 0
        self.temperature = nn.Parameter(torch.tensor(temperature, dtype=torch.float32), requires_grad=False)
        self.scaler = 1 / math.sqrt(d_embed / 2)
        self.proj = nn.Linear(n_dim, d_embed // 2, bias=False)
        self.mlp = nn.Sequential(
            nn.LayerNorm(d_embed),
            nn.Linear(d_embed, d_embed),
            nn.GELU(),
            nn.Linear(d_embed, d_embed),
        )
    def forward(self, x: torch.Tensor):
        """
        Args:
            x (torch.Tensor): [... x D]
        Returns:
            torch.Tensor: [... x E]
        """
        x = self.proj(x * self.temperature)
        x = torch.cat([torch.sin(x), torch.cos(x)], dim=-1) * self.scaler
        return self.mlp(x)